"""
Enhanced FakeTagger System
完整的增强版FakeTagger系统，整合所有5个模块
1. Message Generator
2. Enhanced Image Encoder (with PatchGAN)  
3. SimSwap GAN Simulator
4. Image Decoder
5. Message Decoder
"""

import torch
import torch.nn as nn
import kornia.losses
import os
import numpy as np

# 导入所有模块
from message_generator import MessageGenerator
from enhanced_image_encoder import EnhancedImageEncoder
from simswap_gan_simulator import SimSwapGANSimulator
from decoder import Decoder  # 使用原有的Image Decoder
from message_decoder import MessageDecoder

class EnhancedFakeTaggerSystem(nn.Module):
    """
    完整的增强版FakeTagger系统
    整合5个模块的完整工作流程：
    原图 + 消息生成 -> 嵌入水印 -> 人脸交换 -> 提取消息 -> 解码消息
    """
    
    def __init__(self, message_size=30, redundancy_factor=2, device='cuda', 
                 use_discriminator=True, discriminator_weight=0.1, 
                 simswap_model_path=None, lr=0.001):
        """
        初始化增强版FakeTagger系统
        
        Args:
            message_size: 原始消息长度（默认30，与原项目一致）
            redundancy_factor: 冗余因子（默认2）
            device: 设备类型
            use_discriminator: 是否使用PatchGAN判别器
            discriminator_weight: 判别器权重
            simswap_model_path: SimSwap模型路径
            lr: 学习率
        """
        super(EnhancedFakeTaggerSystem, self).__init__()
        
        self.message_size = message_size
        self.redundancy_factor = redundancy_factor
        self.redundant_message_size = message_size * redundancy_factor
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
        self.use_discriminator = use_discriminator
        
        print(f"Enhanced FakeTagger System initializing...")
        print(f"  - Original message size: {message_size}")
        print(f"  - Redundancy factor: {redundancy_factor}")
        print(f"  - Redundant message size: {self.redundant_message_size}")
        print(f"  - Device: {self.device}")
        print(f"  - Use discriminator: {use_discriminator}")
        
        # 1. Message Generator & Decoder
        print("\n=== Initializing Module 1 & 5: Message Generator & Decoder ===")
        self.message_generator = MessageGenerator(message_size, redundancy_factor)
        self.message_decoder = MessageDecoder(message_size, redundancy_factor)
        
        # 2. Enhanced Image Encoder (with PatchGAN)
        print("\n=== Initializing Module 2: Enhanced Image Encoder ===")
        self.image_encoder = EnhancedImageEncoder(
            message_size=self.redundant_message_size,  # 使用冗余消息长度
            use_discriminator=use_discriminator,
            discriminator_weight=discriminator_weight,
            device=device,
            lr=0.0002  # 编码器使用较小的学习率
        )
        
        # 3. SimSwap GAN Simulator
        print("\n=== Initializing Module 3: SimSwap GAN Simulator ===")
        self.simswap_simulator = SimSwapGANSimulator(
            device=device,
            simswap_model_path=simswap_model_path
        )
        
        # 4. Image Decoder (使用原有的decoder)
        print("\n=== Initializing Module 4: Image Decoder ===")
        self.image_decoder = Decoder(self.redundant_message_size)
        
        # 移动所有模块到设备
        print(f"Moving all modules to device: {self.device}")
        
        # 移动Image Decoder
        self.image_decoder = self.image_decoder.to(self.device)
        
        # 确保Image Encoder也在正确设备上
        self.image_encoder = self.image_encoder.to(self.device)
        
        # 验证设备
        encoder_device = next(self.image_encoder.encoder.parameters()).device
        decoder_device = next(self.image_decoder.parameters()).device
        
        print(f"Image Encoder device: {encoder_device}")
        print(f"Image Decoder device: {decoder_device}")
        
        if self.use_discriminator:
            disc_device = next(self.image_encoder.discriminator.parameters()).device
            print(f"Discriminator device: {disc_device}")
        
        # 确保所有模块都在同一设备上
        if encoder_device != self.device:
            print(f"Warning: Encoder not on correct device!")
        if decoder_device != self.device:
            print(f"Warning: Decoder not on correct device!")
        
        # 系统级损失函数和优化器
        self.criterion_message = nn.MSELoss()
        self.criterion_reconstruction = nn.L1Loss()
        
        # 系统级优化器（训练image_decoder）
        self.opt_image_decoder = torch.optim.Adam(
            self.image_decoder.parameters(),
            lr=lr, betas=(0.5, 0.999)
        )
        
        print(f"\n=== Enhanced FakeTagger System initialized successfully ===")
        print(f"Total modules: 5 (Message Gen/Dec, Enhanced Encoder, SimSwap, Image Decoder)")
    
    def generate_message_batch(self, batch_size):
        """
        生成一批消息（调用模块1）
        
        Args:
            batch_size: 批大小
            
        Returns:
            tuple: (原始消息, 冗余消息)
        """
        return self.message_generator.generate_redundant_message(batch_size, self.device)
    
    def embed_watermark(self, images, redundant_messages):
        """
        在图像中嵌入水印（调用模块2）
        
        Args:
            images: 原始图像
            redundant_messages: 冗余消息
            
        Returns:
            编码后的图像
        """
        return self.image_encoder.encode(images, redundant_messages)
    
    def deepfake_attack(self, encoded_images, swap_type='A'):
        """
        进行深度伪造攻击（调用模块3）
        
        Args:
            encoded_images: 嵌入水印的图像
            swap_type: 交换类型
            
        Returns:
            深度伪造后的图像
        """
        return self.simswap_simulator.deepfake(encoded_images, swap_type)
    
    def extract_messages(self, deepfaked_images):
        """
        从攻击后的图像中提取冗余消息（调用模块4）
        
        Args:
            deepfaked_images: 深度伪造后的图像
            
        Returns:
            提取的冗余消息
        """
        return self.image_decoder(deepfaked_images)
    
    def decode_messages(self, extracted_redundant_messages):
        """
        从冗余消息中解码原始消息（调用模块5）
        
        Args:
            extracted_redundant_messages: 提取的冗余消息
            
        Returns:
            解码的原始消息
        """
        return self.message_decoder.decode_redundant_message(extracted_redundant_messages)
    
    def full_pipeline(self, images, swap_type='A', return_intermediate=True):
        """
        完整的5模块流水线
        
        Args:
            images: 输入图像
            swap_type: 人脸交换类型
            return_intermediate: 是否返回中间结果
            
        Returns:
            dict: 包含各阶段结果的字典
        """
        batch_size = images.shape[0]
        
        # 1. 生成消息
        original_messages, redundant_messages = self.generate_message_batch(batch_size)
        
        # 2. 嵌入水印
        encoded_images = self.embed_watermark(images, redundant_messages)
        
        # 3. 深度伪造攻击
        deepfaked_images = self.deepfake_attack(encoded_images, swap_type)
        
        # 4. 提取冗余消息
        extracted_redundant_messages = self.extract_messages(deepfaked_images)
        
        # 5. 解码原始消息
        decoded_messages = self.decode_messages(extracted_redundant_messages)
        
        results = {
            'original_images': images,
            'original_messages': original_messages,
            'redundant_messages': redundant_messages,
            'encoded_images': encoded_images,
            'deepfaked_images': deepfaked_images,
            'extracted_redundant_messages': extracted_redundant_messages,
            'decoded_messages': decoded_messages
        }
        
        if not return_intermediate:
            # 只返回关键结果
            return {
                'encoded_images': encoded_images,
                'deepfaked_images': deepfaked_images,
                'decoded_messages': decoded_messages,
                'original_messages': original_messages
            }
        
        return results
    
    def train_step(self, images, swap_type='A'):
        """
        完整的训练步骤 - 修复版
        实现端到端训练，改进损失函数
        
        Args:
            images: 输入图像批次
            swap_type: 人脸交换类型
            
        Returns:
            dict: 训练损失
            dict: 训练结果
        """
        batch_size = images.shape[0]
        images = images.to(self.device)
        
        # 生成消息
        original_messages, redundant_messages = self.generate_message_batch(batch_size)
        
        # ==================== 修复：端到端训练 ====================
        
        # === 阶段1: 训练Enhanced Image Encoder (包含PatchGAN对抗训练) ===
        encoder_losses, encoded_images = self.image_encoder.train_step(images, redundant_messages)
        
        # === 阶段2: 进行SimSwap攻击（推理模式，不训练） ===
        with torch.no_grad():
            deepfaked_images = self.deepfake_attack(encoded_images, swap_type)
        
        # === 阶段3: 训练Image Decoder（修复版） ===
        self.opt_image_decoder.zero_grad()
        
        # 从攻击后的图像中提取消息
        extracted_redundant_messages = self.extract_messages(deepfaked_images)
        
        # 修复：使用更合适的损失函数
        
        # 1. 消息重建损失（L1 + L2结合）
        l1_loss = nn.L1Loss()(extracted_redundant_messages, redundant_messages)
        l2_loss = nn.MSELoss()(extracted_redundant_messages, redundant_messages)
        message_extraction_loss = l1_loss + 0.5 * l2_loss
        
        # 2. 二值化损失（鼓励输出接近0/1）
        # 计算sigmoid后的熵损失，鼓励输出接近0或1
        sigmoid_extracted = torch.sigmoid((extracted_redundant_messages - 0.5) * 10)
        binary_loss = -torch.mean(
            sigmoid_extracted * torch.log(sigmoid_extracted + 1e-8) + 
            (1 - sigmoid_extracted) * torch.log(1 - sigmoid_extracted + 1e-8)
        )
        
        # 3. 比特一致性损失
        # 鼓励提取的消息在二值化后与原始消息一致
        binary_original = (redundant_messages > 0.5).float()
        binary_extracted = (extracted_redundant_messages > 0.5).float()
        consistency_loss = nn.MSELoss()(binary_extracted, binary_original)
        
        # 4. 总损失
        total_decoder_loss = (
            message_extraction_loss + 
            0.1 * binary_loss + 
            0.2 * consistency_loss
        )
        
        # 反向传播
        total_decoder_loss.backward()
        
        # 梯度裁剪，防止梯度爆炸
        torch.nn.utils.clip_grad_norm_(self.image_decoder.parameters(), max_norm=1.0)
        
        self.opt_image_decoder.step()
        
        # === 阶段4: 解码消息并计算指标 ===
        with torch.no_grad():
            decoded_messages = self.decode_messages(extracted_redundant_messages)
            
            # 计算指标
            bit_error_rate = self.message_decoder.calculate_bit_error_rate(
                original_messages, decoded_messages
            )
            
            message_accuracy = self.message_decoder.calculate_message_accuracy(
                original_messages, decoded_messages
            )
            
            # 计算软误码率
            soft_ber = self.message_decoder.calculate_soft_bit_error_rate(
                original_messages, decoded_messages
            )
            
            # 计算PSNR（图像质量）
            psnr = -kornia.losses.psnr_loss(encoded_images.detach(), images, 2)
            
            # 额外指标：提取质量
            extraction_quality = nn.MSELoss()(extracted_redundant_messages, redundant_messages)
        
        # 整合所有损失
        all_losses = {
            **encoder_losses,  # 包含编码器和判别器的损失
            'message_extraction_loss': message_extraction_loss.item(),
            'binary_loss': binary_loss.item(),
            'consistency_loss': consistency_loss.item(),
            'total_decoder_loss': total_decoder_loss.item(),
            'bit_error_rate': bit_error_rate,
            'message_accuracy': message_accuracy,
            'soft_ber': soft_ber,
            'psnr': psnr.item(),
            'extraction_quality': extraction_quality.item()
        }
        
        # 训练结果
        train_results = {
            'original_images': images,
            'encoded_images': encoded_images,
            'deepfaked_images': deepfaked_images,
            'original_messages': original_messages,
            'redundant_messages': redundant_messages,
            'extracted_messages': extracted_redundant_messages,
            'decoded_messages': decoded_messages
        }
        
        return all_losses, train_results
    
    def evaluate(self, images, swap_type='A'):
        """
        评估模式（无梯度计算）
        
        Args:
            images: 输入图像
            swap_type: 交换类型
            
        Returns:
            dict: 评估指标
            dict: 评估结果
        """
        self.eval()
        
        with torch.no_grad():
            results = self.full_pipeline(images, swap_type, return_intermediate=True)
            
            # 计算评估指标
            bit_error_rate = self.message_decoder.calculate_bit_error_rate(
                results['original_messages'], results['decoded_messages']
            )
            
            message_accuracy = self.message_decoder.calculate_message_accuracy(
                results['original_messages'], results['decoded_messages']
            )
            
            psnr = -kornia.losses.psnr_loss(
                results['encoded_images'], results['original_images'], 2
            )
            
            metrics = {
                'bit_error_rate': bit_error_rate,
                'message_accuracy': message_accuracy,
                'psnr': psnr.item()
            }
        
        self.train()  # 恢复训练模式
        return metrics, results
    
    def save_system(self, save_dir):
        """
        保存整个系统
        
        Args:
            save_dir: 保存目录
        """
        os.makedirs(save_dir, exist_ok=True)
        
        # 保存各个模块
        self.image_encoder.save_encoder_state(
            os.path.join(save_dir, 'enhanced_image_encoder.pth')
        )
        
        torch.save(
            self.image_decoder.state_dict(),
            os.path.join(save_dir, 'image_decoder.pth')
        )
        
        torch.save(
            self.opt_image_decoder.state_dict(),
            os.path.join(save_dir, 'opt_image_decoder.pth')
        )
        
        # 保存系统配置
        system_config = {
            'message_size': self.message_size,
            'redundancy_factor': self.redundancy_factor,
            'redundant_message_size': self.redundant_message_size,
            'use_discriminator': self.use_discriminator,
            'device': str(self.device)
        }
        
        torch.save(system_config, os.path.join(save_dir, 'system_config.pth'))
        
        # 保存SimSwap信息
        self.simswap_simulator.save_model_info(
            os.path.join(save_dir, 'simswap_info.pth')
        )
        
        print(f"Enhanced FakeTagger System saved to {save_dir}")
    
    def load_system(self, save_dir):
        """
        加载整个系统
        
        Args:
            save_dir: 保存目录
        """
        # 加载系统配置
        system_config = torch.load(
            os.path.join(save_dir, 'system_config.pth'),
            map_location=self.device
        )
        
        print(f"Loading system with config: {system_config}")
        
        # 加载各个模块
        self.image_encoder.load_encoder_state(
            os.path.join(save_dir, 'enhanced_image_encoder.pth')
        )
        
        self.image_decoder.load_state_dict(
            torch.load(
                os.path.join(save_dir, 'image_decoder.pth'),
                map_location=self.device
            )
        )
        
        if os.path.exists(os.path.join(save_dir, 'opt_image_decoder.pth')):
            self.opt_image_decoder.load_state_dict(
                torch.load(
                    os.path.join(save_dir, 'opt_image_decoder.pth'),
                    map_location=self.device
                )
            )
        
        print(f"Enhanced FakeTagger System loaded from {save_dir}")

def test_enhanced_faketagger_system():
    """测试完整的增强版FakeTagger系统"""
    print("=== Testing Enhanced FakeTagger System ===")
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # 创建系统
    system = EnhancedFakeTaggerSystem(
        message_size=30,
        redundancy_factor=2,
        device=device,
        use_discriminator=True,
        discriminator_weight=0.1
    )
    
    # 测试数据
    batch_size = 4
    image_size = 256
    
    test_images = torch.randn(batch_size, 3, image_size, image_size)
    if device == 'cuda':
        test_images = test_images.cuda()
    
    print(f"\nTest data: {test_images.shape}")
    
    # 测试完整流水线
    print("\n--- Testing Full Pipeline ---")
    system.eval()
    metrics, results = system.evaluate(test_images, swap_type='A')
    
    print("Pipeline results:")
    for key, value in results.items():
        if isinstance(value, torch.Tensor):
            print(f"  {key}: {value.shape}")
    
    print("\nEvaluation metrics:")
    for key, value in metrics.items():
        print(f"  {key}: {value:.6f}")
    
    # 测试训练步骤
    print("\n--- Testing Training Step ---")
    system.train()
    losses, train_results = system.train_step(test_images, swap_type='A')
    
    print("Training losses:")
    for key, value in losses.items():
        print(f"  {key}: {value:.6f}")
    
    # 测试不同的swap类型
    print("\n--- Testing Different Swap Types ---")
    for swap_type in ['A', 'B']:
        metrics, _ = system.evaluate(test_images, swap_type=swap_type)
        print(f"Swap type {swap_type} - BER: {metrics['bit_error_rate']:.6f}")
    
    print("\nEnhanced FakeTagger System test completed!")

if __name__ == "__main__":
    test_enhanced_faketagger_system() 